package com.sky.es.util;

import com.alibaba.fastjson.JSONObject;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.client.indices.CreateIndexRequest;
import org.elasticsearch.client.indices.CreateIndexResponse;
import org.elasticsearch.client.indices.GetIndexRequest;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.index.reindex.UpdateByQueryRequest;
import org.elasticsearch.rest.RestStatus;
import org.elasticsearch.script.Script;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.data.elasticsearch.core.ElasticsearchRestTemplate;

import java.io.IOException;
import java.io.Serializable;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class ElasticsearchUtil implements Serializable {

    /**
     * 根据Query和分页条件查询ES库数据
     *
     * @param client ES索引Client对象
     * @param indexName 索引名称
     * @param query 查询条件
     * @param from 开始记录 ES的开始数从0开始
     * @param pageSize 每页大小数
     * @param sortName 排序字段名
     * @param sortOrder 排序规则
     * @return SearchHits查询结果
     */
    public static SearchHits searchByQuery(RestHighLevelClient client, String indexName, QueryBuilder query,
                                           int from, int pageSize, String sortName, SortOrder sortOrder) {
        try {
            SearchSourceBuilder sourceBuilder = SearchSourceBuilder.searchSource();
            sourceBuilder.query(query);
            sourceBuilder.from(from);
            sourceBuilder.size(pageSize);
            sourceBuilder.sort(sortName, sortOrder);
            SearchRequest request = new SearchRequest(indexName);
            request.source(sourceBuilder);
            SearchResponse response = client.search(request, RequestOptions.DEFAULT);
            System.out.println("查询完成编码：" + response.status().getStatus());
            if (response.status().getStatus() == RestStatus.OK.getStatus()) {
                return response.getHits();
            }
            return null;
        } catch (Exception e) {
            return null;
        }
    }

    /**
     * 根据主键ID查询
     * @param client es对象
     * @param indexName 索引名称
     * @param idValue 主键查询值
     * @return json字符串
     * @throws IOException
     */
    public static String findById(RestHighLevelClient client, String indexName, String idValue) throws IOException {
            GetRequest request = new GetRequest(indexName, idValue);
        GetResponse response = client.get(request, RequestOptions.DEFAULT);
        if (response.isExists()) {
            return response.getSourceAsString();
        }
        return null;
    }

    /**
     * 批量添加索引
     *
     * @param client ES索引Client对象
     * @param indexName 索引名
     * @param jsonData json对象集合
     * @param idName 对象主键字段名
     */
    public static boolean batchCreateIndex(RestHighLevelClient client, String indexName, String idName, List<String> jsonData) throws IOException {

        BulkRequest bulkAddRequest = new BulkRequest();

        for (int i = 0; i < jsonData.size(); i++) {
            IndexRequest indexRequest = new IndexRequest(indexName);
            indexRequest.id(JSONObject.parseObject(jsonData.get(i)).getString(idName));
            indexRequest.source(jsonData.get(i), XContentType.JSON);
            bulkAddRequest.add(indexRequest);
        }
        BulkResponse response = client.bulk(bulkAddRequest, RequestOptions.DEFAULT);
        if (response.status().getStatus() == RestStatus.OK.getStatus()) {
            return true;
        }
        return false;
    }

    /**
     * 添加单条索引记录
     *
     * @param client ES索引Client对象
     * @param indexName 索引名
     * @param jsonData 索引对象的json串
     * @param idName 记录对应的主键名
     */
    public static boolean addEsIndex(RestHighLevelClient client, String indexName, String idName, String jsonData) throws IOException {
        IndexRequest request = new IndexRequest(indexName);
        request.source(jsonData, XContentType.JSON);
        request.id(JSONObject.parseObject(jsonData).getString(idName));
        IndexResponse indexResponse = client.index(request, RequestOptions.DEFAULT);
        if (indexResponse != null && indexResponse.status().getStatus() == RestStatus.CREATED.getStatus()) {
            return true;
        }
        return false;
    }

    /**
     * 更新ES记录
     *
     * @param client ES索引Client对象
     * @param indexName 索引名
     * @param idName 记录id名
     * @param jsonData 要更新的记录json串
     * @return true:更新成功。false：失败
     */
    public static boolean updateIndex(RestHighLevelClient client, String indexName, String idName, String jsonData) throws IOException {
        JSONObject jsonObject = JSONObject.parseObject(jsonData);
        // 1.记录不存在，会报DocumentMissingException
        UpdateRequest updateRequest = new UpdateRequest();
        // 设置更新索引名
        updateRequest.index(indexName);
        // 更新记录的主键ID
        updateRequest.id(jsonObject.getString(idName));
        // 需要更新的记录
        updateRequest.doc(jsonData, XContentType.JSON);
        UpdateResponse response = client.update(updateRequest, RequestOptions.DEFAULT);
        if (response.status().getStatus() == RestStatus.OK.getStatus()) {
            return true;
        }
        return false;
    }

    /**
     * 查询某个索引下的记录总数
     *
     * @param client ES索引Client对象
     * @param indexName 索引名称
     * @return
     */
    public static Long getCount(RestHighLevelClient client, String indexName) throws IOException {
        Long count = 0L;
        CountRequest request = new CountRequest(indexName);
        CountResponse response = client.count(request, RequestOptions.DEFAULT);
        if (response.status().getStatus() == RestStatus.OK.getStatus()) {
            return response.getCount();
        }
        return count;
    }

    /**
     * 创建索引
     *
     * @param elasticsearchRestTemplate ES索模板对象
     * @param clazz ES mapping 实体类
     * @return
     */
    public static boolean createIndexAndMapping(ElasticsearchRestTemplate elasticsearchRestTemplate, Class clazz) {
        boolean create = elasticsearchRestTemplate.createIndex(clazz);
        boolean mapping = elasticsearchRestTemplate.putMapping(clazz);
        if (create && mapping) {
            return true;
        }
        return false;
    }


    /**
     * 删除单条索引,根据指定ID删除文档
     *
     * @param client ES索引Client对象
     * @param indexName 索引名
     * @param idValue 文档ID值
     */
    public static boolean deleteIndex(RestHighLevelClient client, String indexName, String idValue) throws IOException {
        DeleteRequest request = new DeleteRequest(indexName);
        request.id(idValue);
        DeleteResponse response = client.delete(request, RequestOptions.DEFAULT);
        if (response.status().getStatus() == RestStatus.OK.getStatus()) {
            return true;
        }
        return false;
    }

    /**
     * 条件删除
     * @param client es对象
     * @param indexName 索引名称
     * @param query 条件
     * @return
     * @throws IOException
     */
    public static boolean deleteIndexByQuery(RestHighLevelClient client, String indexName, QueryBuilder query) throws IOException {
        DeleteByQueryRequest request = new DeleteByQueryRequest(indexName);
        request.setQuery(query);
        request.setConflicts("proceed");
        request.setRefresh(true);
        BulkByScrollResponse response = client.deleteByQuery(request, RequestOptions.DEFAULT);
        if (response.getTotal() > 0) {
            return true;
        }
        return false;
    }

    /**
     * 条件修改
     * 参考文档：https://segmentfault.com/a/1190000017123753?utm_source=tag-newest
     *         https://www.elastic.co/guide/en/elasticsearch/client/java-api/current/java-docs-update-by-query.html
     * @param client es对象
     * @param indexName 索引名称
     * @param query 条件
     * @param script 脚本
     * @return
     * @throws IOException
     */
    public static boolean updateIndexByQuery(RestHighLevelClient client, String indexName, QueryBuilder query, Script script) throws IOException {
        UpdateByQueryRequest request = new UpdateByQueryRequest(indexName);
        request.setQuery(query);
        request.setConflicts("proceed");
        request.setRefresh(true);
        request.setScript(script);
        BulkByScrollResponse response = client.updateByQuery(request, RequestOptions.DEFAULT);
        if (response.getTotal() > 0) {
            return true;
        }
        return false;
    }

    /**
     * 判断指定的索引名是否存在
     *
     * @param client ES索引Client对象
     * @param indexName 索引名
     * @return 存在：true; 不存在：false;
     */
    public static boolean isExistsIndex(RestHighLevelClient client, String indexName) {
        boolean exists = false;
        try {
            exists = client.indices().exists(new GetIndexRequest(indexName), RequestOptions.DEFAULT);
        } catch (IOException e) {
            e.printStackTrace();
        }
        System.out.println("existsIndex: " + exists);
        return exists;
    }



    /**
     * 创建索引, 以及索引Mapping
     *
     * @param client RestHighLevelClient 对象
     * @param indexName 索引名称
     * @param builder 索引mapping
     * @return
     */
    public static boolean createIndex(RestHighLevelClient client, String indexName, XContentBuilder builder) throws IOException {
        Map<String, Integer> map = new HashMap<String, Integer>();
        //是数据分片数，如果这里不设置map则为1个，如果使用@Document注解默认则为5个
        map.put("number_of_shards", 4);
        //是数据备份数，如果只有一台机器，设置为0
        map.put("number_of_replicas", 1);
        CreateIndexRequest request = new CreateIndexRequest(indexName);
        request.settings(map);
        request.mapping(builder);
        CreateIndexResponse createIndexResponse = client.indices().create(request, RequestOptions.DEFAULT);
        return createIndexResponse.isAcknowledged();
    }
}
